import gradio as gr
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import io
import torch
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
doctor_desc = [
    '正常窦性心律', '窦性心动过速', '窦性心动过缓', '房性早搏', '室性早搏',
    '心房颤动', '心房扑动', '阵发性室上性心动过速', '室性心动过速', '心室颤动',
    '左心房肥大/右心房肥大', '左心室肥厚/右心室肥厚', '左束支传导阻滞/右束支传导阻滞', '心肌缺血改变', '急性心肌梗死/陈旧性心肌梗死',
    '急性心包炎', '高钾血症/低钾血症/高钙血症/低钙血症', 'WPW综合征', '二度窦房传导阻滞', '一度房室传导阻滞/二度房室传导阻滞/三度房室传导阻滞',
    'QT间期延长', 'Brugada综合征心电图表现', '早期复极化改变', '肺源性P波', 'QRS低电压'
]

cn_text = [
    '窦性心律', '窦性心动过速', '窦性心动过缓', '房性早搏', '室性早搏',
    '房颤', '房扑', '室上性心动过速', '室性心动过速', '心室颤动',
    '心房肥大', '心室肥大', '束支传导阻滞', '心肌缺血', '心肌梗死',
    '心包炎', '电解质紊乱', '预激综合征', '窦房传导阻滞', '房室传导阻滞',
    'QT间期延长', 'Brugada综合征', '早期复极化', '肺性P波', '低电压'
]

text = [
    'Sinus rhythm', 'Sinus tachycardia', 'Sinus bradycardia', 'Atrial premature complex', 'Ventricular premature complex',
    'Atrial fibrillation', 'Atrial flutter', 'Supraventricular tachycardia', 'Ventricular tachycardia', 'Ventricular fibrillation',
    'Atrial enlargement', 'Ventricular hypertrophy', 'Bundle branch block', 'Myocardial ischemia', 'Myocardial infarction',
    'Pericarditis', 'Electrolyte disturbances', 'Pre-excitation syndrome', 'Sinoatrial block', 'Atrioventricular block',
    'QT interval prolongation', 'Brugada syndrome', 'Early repolarization', 'P pulmonale', 'Low voltage'
]

module_path = "openai/clip-vit-base-patch32"

# 加载 Transformer 模型
def load_model():
    model = CLIPModel.from_pretrained(module_path)
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    print(device)
    model.to(device)
    return model, processor

def ecg_classifier(img_file):
    if img_file is None:
        return "请上传图片"
    if not (img_file.name.endswith('.jpg') or img_file.name.endswith('.jpeg') or img_file.name.endswith('.png')):
        return "图片格式不正确"
    model, processor = load_model()
    print("加载模型成功")
    # import requests
    # # 下载模型
    # API_URL = "https://ai.gitee.com/api/endpoints/21jhf/clip-vit-large-patch14-336-8419/inference"
    # headers = {
    #     "Authorization": "Bearer eyJpc3MiOiJodHRwczovL2FpLmdpdGVlLmNvbSIsInN1YiI6IjM1NjU0In0.98IHoSuCTLDnwAU7NEPP95VQei1t-lid-b9raGE3quK430kuo0Qw1DqnKHvCl4RWM6jHD-wm5IARa3JWkNFtBQ",
    #     "Content-Type": "image/jpeg"
    # }
    #
    # def query(filename):
    #     with open(filename, "rb") as f:
    #         data = f.read()
    #     response = requests.post(API_URL, headers=headers, data=data)
    #     return response.json()
    #
    # output = query("ecg1.jpg")
    with open(img_file, 'rb') as f:
        img_data = f.read()

    img_pil = Image.open(io.BytesIO(img_data))
    width, height = img_pil.size
    format = img_pil.format

    # 构造一些基于图片的信息
    info = f"文件名: {img_file.name}\n尺寸: {width}x{height}像素\n格式: {format}"
    print(img_file)
    print("开始识别")
    # 预处理图像
    img_array = np.array(img_pil)
    image_inputs = processor(images=img_array, padding=True, return_tensors="pt")
    # 这里解决使用GPU后模型权重和输入数据在不同的设备上的问题，将处理后的输入移动到正确的设备
    image_inputs = {k: v.to(device) for k, v in image_inputs.items()}
    # 预处理所有文本
    text_inputs = processor(text=text, padding=True, truncation=True, return_tensors="pt")
    # 这里解决使用GPU后模型权重和输入数据在不同的设备上的问题，将处理后的输入移动到正确的设备
    text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
    # 使用 torch.no_grad() 来提高推理速度
    with torch.no_grad():
        # 获取图像特征
        image_features = model.get_image_features(**image_inputs)

        # 获取文本特征
        text_features = model.get_text_features(**text_inputs)

        # 计算相似度得分
        logits = (image_features @ text_features.T).squeeze(0)

        # 使用 softmax 将得分转换为概率分布
        probs = F.softmax(logits, dim=-1)

    # 将结果转换为列表
    probs_list = probs.tolist()

    for i, (t, p) in enumerate(zip(text, probs_list)):
        result_line = f"识别结果: {t}->{doctor_desc[i]} 概率: {p:.4f}\n"  # 格式化每行结果，保留4位小数
        info += result_line  # 将每行结果添加到 info 字符串
        print(result_line, end='')

        # 返回信息作为文本
    return info


ecg = gr.Interface(fn=ecg_classifier,
                   inputs=gr.File(label="上传图片", file_types=["jpg", "jpeg", "png"]),
                   outputs=gr.Textbox(lines=10, label="识别结论"),
                   live=True,  # 启用实时更新
                   title="心电图识别系统-赢优科技有限公司"  # 界面标题
                   )
ecg.launch(server_name='0.0.0.0')